{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "38c1a328-fd86-4c5f-bd54-b8664f433608",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Simple retrieval augmented generation with OpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32f8484d-2e35-472a-9b24-1a30ec1d144b",
   "metadata": {},
   "source": [
    "<!-- TABS -->\n",
    "## Connect to superduper"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06d66021-ce62-4021-a2c5-158dee92b3bb",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    ":::note\n",
    "Note that this is only relevant if you are running superduper in development mode.\n",
    "Otherwise refer to \"Configuring your production system\".\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3ef70f6d-a189-460a-8864-241a689624e2",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "APPLY = True\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ['SUPERDUPER_SECRETS_VOLUME'] = '~/data/secrets/snowflake_dev'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cb029a5e-fedf-4f07-8a31-d220cfbfbb3d",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building Datalayer...\n",
      "Building Datalayer... DONE\n"
     ]
    }
   ],
   "source": [
    "from superduper import superduper, CFG\n",
    "import os\n",
    "\n",
    "db = superduper('mongomock://', initialize_cluster=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4e7902bd",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import requests\n",
    "import io\n",
    "from superduper import logging\n",
    "\n",
    "\n",
    "def getter():\n",
    "    logging.info('Downloading data...')\n",
    "    response = requests.get('https://superduperdb-public-demo.s3.amazonaws.com/text.json')\n",
    "    logging.info('Downloading data... (Done)')\n",
    "    data = json.loads(response.content.decode('utf-8'))\n",
    "    return [{'x': r} for r in data]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1ef8dd07-1b47-4dce-84dd-a081d1f5ee9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "if APPLY:\n",
    "    data = getter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "14067b23-7958-4c93-9311-2525b133bfca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from superduper import Table\n",
    "\n",
    "table = Table('<var:table_name>', fields={'x': 'str'})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4885060-450c-4446-b5fb-e58591c0d1ff",
   "metadata": {},
   "source": [
    "Create plugin:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f879a4df-6f60-4c80-8fdd-285a1ef92c90",
   "metadata": {},
   "outputs": [],
   "source": [
    "from superduper import Plugin\n",
    "\n",
    "plugin = Plugin(path='./rag_plugin.py')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54fea927-ee4a-44cd-aaf2-634b574c316d",
   "metadata": {},
   "source": [
    "<!-- TABS -->\n",
    "## Apply a chunker for search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06d90bda-e8c4-494e-a38c-837fb63689ae",
   "metadata": {},
   "source": [
    ":::note\n",
    "Note that applying a chunker is ***not*** mandatory for search.\n",
    "If your data is already chunked (e.g. short text snippets or audio) or if you\n",
    "are searching through something like images, which can't be chunked, then this\n",
    "won't be necessary.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "93d21872-d4dc-40dc-abab-fb07ba102ea3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from superduper import Listener\n",
    "from rag_plugin import Chunker\n",
    "\n",
    "upstream_listener = Listener(\n",
    "    model=Chunker(identifier='chunker'),\n",
    "    select=db['<var:table_name>'],\n",
    "    key='x',\n",
    "    identifier='chunker',\n",
    "    flatten=True,\n",
    "    upstream=[plugin],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9a2cd87-723f-4cee-87c7-9b8181c9e54b",
   "metadata": {},
   "source": [
    "<!-- TABS -->\n",
    "## Build text embedding model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10753ea4-9893-4056-813d-7d6ddf78ce02",
   "metadata": {},
   "source": [
    "OpenAI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a9b1f538-65ca-499e-b6d0-2dd733f81723",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from superduper_openai import OpenAIEmbedding\n",
    "\n",
    "openai_embedding = OpenAIEmbedding(\n",
    "    identifier='text-embedding',\n",
    "    model='<var:embedding_model>',\n",
    "    datatype='vector[float32:1536]',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f31843db-8638-458a-a770-96a79041be88",
   "metadata": {},
   "source": [
    "## Create vector-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4663fa4b-c2ec-427d-bf8b-b8b109cc2ccf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from superduper import VectorIndex, Listener\n",
    "\n",
    "vector_index_name = 'vectorindex'\n",
    "\n",
    "vector_index = VectorIndex(\n",
    "    vector_index_name,\n",
    "    indexing_listener=Listener(\n",
    "        key=upstream_listener.outputs,\n",
    "        select=db[upstream_listener.outputs],\n",
    "        model=openai_embedding,\n",
    "        identifier='embeddinglistener',\n",
    "        upstream=[upstream_listener],\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1179a67b-4e40-496b-9851-98f32d42faa0",
   "metadata": {},
   "source": [
    "<!-- TABS -->\n",
    "## Build LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "75faf501-f0cf-4707-a165-5a05cfb14bbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from superduper_openai import OpenAIChatCompletion\n",
    "\n",
    "\n",
    "llm_openai = OpenAIChatCompletion(\n",
    "    identifier='llm-model',\n",
    "    model='<var:llm_model>',\n",
    "    datatype='str',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60ae6203-dcc4-493c-a8f8-f727f0f75778",
   "metadata": {},
   "source": [
    "## Answer question with LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "44baeb09-6f35-4cf2-b814-46283a59f7e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rag_plugin import RAGModel\n",
    "\n",
    "\n",
    "prompt_template = (\n",
    "    \"Use the following context snippets, these snippets are not ordered!, Answer the question based on this context.\\n\"\n",
    "    \"These snippets are samples from our internal data-repositories, and should be used exclusively and as a matter\"\n",
    "    \" of priority to answer the question. Please answer in 20 words or less.\\n\"\n",
    "    \"{context}\\n\"\n",
    "    \"Here is the question: {query}\"\n",
    ")\n",
    "\n",
    "\n",
    "rag = RAGModel(\n",
    "    'simple_rag',\n",
    "    select=db[upstream_listener.outputs].select().like({upstream_listener.outputs: '<var:query>'}, vector_index=vector_index_name, n=5),\n",
    "    prompt_template=prompt_template,\n",
    "    key=upstream_listener.outputs,\n",
    "    llm=llm_openai,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "183bf5b6-4644-4e4c-b65b-e6bafdc6b49f",
   "metadata": {},
   "source": [
    "By applying the RAG model to the database, it will subsequently be accessible for use in other services."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c974643b-e642-40ea-942f-4d90e0d1bbe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from superduper import Streamlit, Plugin\n",
    "from rag_plugin import demo_func\n",
    "\n",
    "demo = Streamlit('simple-rag-demo', demo_func=demo_func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e6787c78-4b14-4a72-818b-450408a74331",
   "metadata": {},
   "outputs": [],
   "source": [
    "from superduper import Application\n",
    "\n",
    "app = Application(\n",
    "    'simple-rag-app',\n",
    "    upstream=[table, plugin],\n",
    "    components=[\n",
    "        upstream_listener,\n",
    "        vector_index,\n",
    "        rag,\n",
    "        demo,\n",
    "    ],\n",
    "    variables={\n",
    "        'table_name': 'docs',\n",
    "        'id_field': '_id',\n",
    "        'embedding_model': 'text-embedding-ada-002',\n",
    "        'llm_model': 'gpt-3.5-turbo',\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0770487c-9e26-4bb8-8554-37fc4a8ca24c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Application:simple-rag-app:6c4f73b9d964d6faeb\n",
       "├── upstream\n",
       "│   ├── [0] Table:docs:3b21dc32e8beebd92b\n",
       "│   │   ├── fields\n",
       "│   │   │   └── x: str\n",
       "│   │   └── primary_id: id\n",
       "│   └── [1] Plugin:plugin-rag_plugin_py:cbcdd2255fcf92ea1a\n",
       "│       ├── path: /Users/dodo/.superduper/plugins/8f5f4060b2f9deb999/rag_plugin.py\n",
       "│       └── cache_path: ~/.superduper/plugins\n",
       "├── components\n",
       "│   ├── [0] Listener:chunker:1862a8fc1629f85fe6\n",
       "│   │   ├── upstream\n",
       "│   │   │   └── [0] Plugin:plugin-rag_plugin_py:cbcdd2255fcf92ea1a\n",
       "│   │   │       ├── path: /Users/dodo/.superduper/plugins/8f5f4060b2f9deb999/rag_plugin.py\n",
       "│   │   │       └── cache_path: ~/.superduper/plugins\n",
       "│   │   ├── cdc_table: docs\n",
       "│   │   ├── key: x\n",
       "│   │   ├── model: Chunker:chunker:b12cd565989b72fc4c\n",
       "│   │   │   ├── datatype: str\n",
       "│   │   │   └── chunk_size: 200\n",
       "│   │   ├── select: docs\n",
       "│   │   └── flatten: True\n",
       "│   ├── [1] VectorIndex:vectorindex:851e610e9d2ace0b54\n",
       "│   │   ├── cdc_table: _outputs__embeddinglistener__ac017bf2de7186ffcd\n",
       "│   │   ├── indexing_listener: Listener:embeddinglistener:a67199bb7c0c1ac162\n",
       "│   │   │   ├── upstream\n",
       "│   │   │   │   └── [0] Listener:chunker:1862a8fc1629f85fe6\n",
       "│   │   │   │       ├── upstream\n",
       "│   │   │   │       │   └── [0] Plugin:plugin-rag_plugin_py:81c4a51cf8d7441211\n",
       "│   │   │   │       │       ├── path: /Users/dodo/.superduper/plugins/ecb6f193ac13ae7cd5/rag_plugin.py\n",
       "│   │   │   │       │       └── cache_path: ~/.superduper/plugins\n",
       "│   │   │   │       ├── cdc_table: docs\n",
       "│   │   │   │       ├── key: x\n",
       "│   │   │   │       ├── model: Chunker:chunker:b12cd565989b72fc4c\n",
       "│   │   │   │       │   ├── datatype: str\n",
       "│   │   │   │       │   └── chunk_size: 200\n",
       "│   │   │   │       ├── select: docs\n",
       "│   │   │   │       └── flatten: True\n",
       "│   │   │   ├── cdc_table: _outputs__chunker__1862a8fc1629f85fe6\n",
       "│   │   │   ├── key: _outputs__chunker__1862a8fc1629f85fe6\n",
       "│   │   │   ├── model: OpenAIEmbedding:text-embedding:2b907268839faa545d\n",
       "│   │   │   │   ├── datatype: vector\n",
       "│   │   │   │   ├── model: text-embedding-ada-002\n",
       "│   │   │   │   ├── max_batch_size: 8\n",
       "│   │   │   │   └── batch_size: 100\n",
       "│   │   │   └── select: _outputs__chunker__1862a8fc1629f85fe6\n",
       "│   │   └── measure: cosine\n",
       "│   ├── [2] RAGModel:simple_rag:233661fe822e51dd90\n",
       "│   │   ├── datatype: str\n",
       "│   │   ├── prompt_template: Use the following context snippets, these snippets are not ordered!, Answer the \n",
       "│   │   │   question based on this context.\n",
       "│   │   │   These snippets are samples from our internal data-repositories, and should be used exclusively and as a\n",
       "│   │   │   matter of priority to answer the question. Please answer in 20 words or less.\n",
       "│   │   │   {context}\n",
       "│   │   │   Here is the question: {query}\n",
       "│   │   ├── select: _outputs__chunker__1862a8fc1629f85fe6.select().like({'_outputs__chunker__1862a8fc1629f85fe6': \n",
       "│   │   │   '&lt;var:query&gt;'}, \"vectorindex\", n=5)\n",
       "│   │   ├── key: _outputs__chunker__1862a8fc1629f85fe6\n",
       "│   │   └── llm: OpenAIChatCompletion:llm-model:ffce2ddfa3a6a5288f\n",
       "│   │       ├── datatype: str\n",
       "│   │       ├── model: gpt-3.5-turbo\n",
       "│   │       ├── max_batch_size: 8\n",
       "│   │       └── batch_size: 1\n",
       "│   └── [3] Streamlit:simple-rag-demo:b33aaf6fa76d173336\n",
       "│       └── demo_func: &lt;function demo_func at 0x1439de3b0&gt;\n",
       "└── variables\n",
       "    ├── table_name: docs\n",
       "    ├── id_field: _id\n",
       "    ├── embedding_model: text-embedding-ada-002\n",
       "    └── llm_model: gpt-3.5-turbo\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Application:simple-rag-app:6c4f73b9d964d6faeb\n",
       "├── upstream\n",
       "│   ├── [0] Table:docs:3b21dc32e8beebd92b\n",
       "│   │   ├── fields\n",
       "│   │   │   └── x: str\n",
       "│   │   └── primary_id: id\n",
       "│   └── [1] Plugin:plugin-rag_plugin_py:cbcdd2255fcf92ea1a\n",
       "│       ├── path: /Users/dodo/.superduper/plugins/8f5f4060b2f9deb999/rag_plugin.py\n",
       "│       └── cache_path: ~/.superduper/plugins\n",
       "├── components\n",
       "│   ├── [0] Listener:chunker:1862a8fc1629f85fe6\n",
       "│   │   ├── upstream\n",
       "│   │   │   └── [0] Plugin:plugin-rag_plugin_py:cbcdd2255fcf92ea1a\n",
       "│   │   │       ├── path: /Users/dodo/.superduper/plugins/8f5f4060b2f9deb999/rag_plugin.py\n",
       "│   │   │       └── cache_path: ~/.superduper/plugins\n",
       "│   │   ├── cdc_table: docs\n",
       "│   │   ├── key: x\n",
       "│   │   ├── model: Chunker:chunker:b12cd565989b72fc4c\n",
       "│   │   │   ├── datatype: str\n",
       "│   │   │   └── chunk_size: 200\n",
       "│   │   ├── select: docs\n",
       "│   │   └── flatten: True\n",
       "│   ├── [1] VectorIndex:vectorindex:851e610e9d2ace0b54\n",
       "│   │   ├── cdc_table: _outputs__embeddinglistener__ac017bf2de7186ffcd\n",
       "│   │   ├── indexing_listener: Listener:embeddinglistener:a67199bb7c0c1ac162\n",
       "│   │   │   ├── upstream\n",
       "│   │   │   │   └── [0] Listener:chunker:1862a8fc1629f85fe6\n",
       "│   │   │   │       ├── upstream\n",
       "│   │   │   │       │   └── [0] Plugin:plugin-rag_plugin_py:81c4a51cf8d7441211\n",
       "│   │   │   │       │       ├── path: /Users/dodo/.superduper/plugins/ecb6f193ac13ae7cd5/rag_plugin.py\n",
       "│   │   │   │       │       └── cache_path: ~/.superduper/plugins\n",
       "│   │   │   │       ├── cdc_table: docs\n",
       "│   │   │   │       ├── key: x\n",
       "│   │   │   │       ├── model: Chunker:chunker:b12cd565989b72fc4c\n",
       "│   │   │   │       │   ├── datatype: str\n",
       "│   │   │   │       │   └── chunk_size: 200\n",
       "│   │   │   │       ├── select: docs\n",
       "│   │   │   │       └── flatten: True\n",
       "│   │   │   ├── cdc_table: _outputs__chunker__1862a8fc1629f85fe6\n",
       "│   │   │   ├── key: _outputs__chunker__1862a8fc1629f85fe6\n",
       "│   │   │   ├── model: OpenAIEmbedding:text-embedding:2b907268839faa545d\n",
       "│   │   │   │   ├── datatype: vector\n",
       "│   │   │   │   ├── model: text-embedding-ada-002\n",
       "│   │   │   │   ├── max_batch_size: 8\n",
       "│   │   │   │   └── batch_size: 100\n",
       "│   │   │   └── select: _outputs__chunker__1862a8fc1629f85fe6\n",
       "│   │   └── measure: cosine\n",
       "│   ├── [2] RAGModel:simple_rag:233661fe822e51dd90\n",
       "│   │   ├── datatype: str\n",
       "│   │   ├── prompt_template: Use the following context snippets, these snippets are not ordered!, Answer the \n",
       "│   │   │   question based on this context.\n",
       "│   │   │   These snippets are samples from our internal data-repositories, and should be used exclusively and as a\n",
       "│   │   │   matter of priority to answer the question. Please answer in 20 words or less.\n",
       "│   │   │   {context}\n",
       "│   │   │   Here is the question: {query}\n",
       "│   │   ├── select: _outputs__chunker__1862a8fc1629f85fe6.select().like({'_outputs__chunker__1862a8fc1629f85fe6': \n",
       "│   │   │   '<var:query>'}, \"vectorindex\", n=5)\n",
       "│   │   ├── key: _outputs__chunker__1862a8fc1629f85fe6\n",
       "│   │   └── llm: OpenAIChatCompletion:llm-model:ffce2ddfa3a6a5288f\n",
       "│   │       ├── datatype: str\n",
       "│   │       ├── model: gpt-3.5-turbo\n",
       "│   │       ├── max_batch_size: 8\n",
       "│   │       └── batch_size: 1\n",
       "│   └── [3] Streamlit:simple-rag-demo:b33aaf6fa76d173336\n",
       "│       └── demo_func: <function demo_func at 0x1439de3b0>\n",
       "└── variables\n",
       "    ├── table_name: docs\n",
       "    ├── id_field: _id\n",
       "    ├── embedding_model: text-embedding-ada-002\n",
       "    └── llm_model: gpt-3.5-turbo\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "app.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "137448de-eeee-4aa1-9202-bf26aca68aee",
   "metadata": {},
   "outputs": [],
   "source": [
    "if APPLY:\n",
    "    db.apply(app, force=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "573cee13-f0be-4cb0-8f27-7fbebdceeb3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "if APPLY:\n",
    "    db['docs'].insert(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f416ebb-cba2-4c00-b4ba-8b25cc5a13f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "db['docs'].execute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "990aa4a3-7674-466f-8efe-43897c06c03a",
   "metadata": {},
   "outputs": [],
   "source": [
    "app.export('.', format='json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
